import asyncio
import json
import os
import sys
from typing import List

from mcp import ClientSession
from mcp.client.sse import sse_client
from openai import AsyncOpenAI
from dotenv import load_dotenv, find_dotenv
import aiohttp
class MCPClient:
    def __init__(self, model_name: str, base_url: str, api_key: str, server_urls: List[str]):
        """
        初始化 MCP 客户端，连接 OpenAI 接口。

        :param model_name: 使用的模型名称，例如 "deepseek-chat"。
        :param base_url: OpenAI 接口的基础地址，例如 "https://api.deepseek.com/v1"。
        :param api_key: OpenAI API 密钥，用于身份验证。
        :param server_urls: SSE 服务地址列表，用于连接多个服务器。
        """
        self.model_name = model_name
        self.server_urls = server_urls
        self.upload_url = "http://192.168.0.184:6001"
        self.sessions = {}  # 存储每个服务器的会话及其上下文：server_id -> (session, session_context, streams_context)
        self.tool_mapping = {}  # 工具映射：prefixed_name -> (session, original_tool_name)

        # 初始化 OpenAI 异步客户端
        self.client = AsyncOpenAI(base_url=base_url, api_key=api_key)
    async def initialize_sessions(self):
        """
        初始化与所有 SSE 服务器的连接，并获取可用工具列表。
        """
        for i, server_url in enumerate(self.server_urls):
            server_id = f"server{i}"  # 为每个服务器生成唯一标识符
            # 创建 SSE 客户端并进入上下文
            streams_context = sse_client(url=server_url)
            streams = await streams_context.__aenter__()
            session_context = ClientSession(*streams)
            session = await session_context.__aenter__()
            await session.initialize()

            # 存储会话及其上下文
            self.sessions[server_id] = (session, session_context, streams_context)

            # 获取工具列表并建立映射
            response = await session.list_tools()
            for tool in response.tools:
                prefixed_name = f"{server_id}_{tool.name}"  # 为工具名添加服务器前缀
                self.tool_mapping[prefixed_name] = (session, tool.name)
            print(f"已连接到 {server_url}，工具列表：{[tool.name for tool in response.tools]}")

    async def cleanup(self):
        """
        清理所有会话和连接资源，确保无资源泄漏。
        """
        for server_id, (session, session_context, streams_context) in self.sessions.items():
            await session_context.__aexit__(None, None, None)  # 退出会话上下文
            await streams_context.__aexit__(None, None, None)  # 退出 SSE 流上下文
        print("所有会话已清理。")

    async def process_query(self, query: str) -> str:
        """
        处理用户的自然语言查询，通过工具调用完成任务并返回结果。

        :param query: 用户输入的查询字符串。
        :return: 处理后的回复文本。
        """
        messages = [{"role": "user", "content": query}]  # 初始化消息列表

        # 收集所有可用工具
        available_tools = []
        for server_id, (session, _, _) in self.sessions.items():
            response = await session.list_tools()
            for tool in response.tools:
                prefixed_name = f"{server_id}_{tool.name}"
                available_tools.append({
                    "type": "function",
                    "function": {
                        "name": prefixed_name,
                        "description": tool.description,
                        "parameters": tool.inputSchema,
                    },
                })
        print('all tools is ',available_tools)
        # 确保传递的消息是 UTF-8 编码
        print('msg is ',messages)
        print('model name is ',self.model_name)
        # 向模型发送初始请求
        try:
            response = await self.client.chat.completions.create(
                model=self.model_name,
                # messages=messages,
                messages= messages,
                tools=available_tools,
            )
        except Exception as e:
            print('error is ',e)
        print("response is ",response)
        final_text = []  # 存储最终回复内容
        message = response.choices[0].message
        final_text.append(message.content or "")  # 添加模型的初始回复

        # 处理工具调用
        while message.tool_calls:
            for tool_call in message.tool_calls:
                prefixed_name = tool_call.function.name
                if prefixed_name in self.tool_mapping:
                    session, original_tool_name = self.tool_mapping[prefixed_name]
                    tool_args = json.loads(tool_call.function.arguments)
                    try:
                        result = await session.call_tool(original_tool_name, tool_args)
                    except Exception as e:
                        result = {"content": f"调用工具 {original_tool_name} 出错：{str(e)}"}
                        print(result["content"])
                    final_text.append(f"[调用工具 {prefixed_name} 参数: {tool_args}]")
                    final_text.append(f"工具结果: {result.content}")
                    messages.extend([
                        {
                            "role": "assistant",
                            "tool_calls": [{
                                "id": tool_call.id,
                                "type": "function",
                                "function": {"name": prefixed_name, "arguments": json.dumps(tool_args)},
                            }],
                        },
                        {"role": "tool", "tool_call_id": tool_call.id, "content": str(result.content)},
                    ])
                else:
                    print(f"工具 {prefixed_name} 未找到")
                    final_text.append(f"工具 {prefixed_name} 未找到")

            # 获取工具调用后的后续回复
            response = await self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                tools=available_tools,
            )
            message = response.choices[0].message
            if message.content:
                final_text.append(message.content)

        return "\n".join(final_text)

    # async def chat_loop(self):
    #     """
    #     启动命令行交互式对话循环，接受用户输入并显示回复。
    #     """
    #     print("\nMCP 客户端已启动，输入你的问题，输入 'quit' 退出。")
    #     while True:
    #         try:
    #             query = input("\n问题: ").strip()
    #             if query.lower() == "quit":
    #                 break
    #             response = await self.process_query(query)
    #             print("\n" + response)
    #         except Exception as e:
    #             print(f"\n发生错误: {str(e)}")
    async def handle_file_upload(self):
        """处理文件上传流程"""
        file_path = input("请输入要上传的文件路径: ").strip()
        if not os.path.exists(file_path):
            print(f"文件不存在: {file_path}")
            return

        try:
            # 显示上传进度
            print(f"\n正在上传文件: {os.path.basename(file_path)}...")
            
            # 调用上传方法
            upload_response = await self.upload_file(file_path)
            
            if upload_response.get("status") == "processing":
                file_id = upload_response["file_id"]
                print(f"文件已接收，处理中 (ID: {file_id})")
                await self.monitor_progress(file_id)
            else:
                print("上传失败:", upload_response.get("error", "未知错误"))
                
        except Exception as e:
            print(f"上传过程中出错: {str(e)}")

    async def upload_file(self, file_path):
        """实际执行文件上传"""
        try:
            with open(file_path, "rb") as f:
                # 明确构造FormData对象
                data = aiohttp.FormData()
                data.add_field(
                    "file",
                    f,
                    filename=os.path.basename(file_path),
                    content_type="application/octet-stream"
                )
                
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        f"{self.upload_url}/upload",
                        data=data
                    ) as response:
                        return await response.json()
        except Exception as e:
            return {"error": str(e)}

    async def monitor_progress(self, file_id):
        """监听文件处理进度"""
        print("等待处理完成... (按Ctrl+C中断)")
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    f"{self.upload_url}/progress/{file_id}"
                ) as response:
                    async for line in response.content:
                        if line:
                            decoded = line.decode("utf-8").strip()
                            if decoded.startswith("data:"):
                                progress = decoded[5:].strip()
                                print(f"\r处理进度: {progress}", end="", flush=True)
                                if progress == "处理完成":
                                    print("\n文件处理完成!")
                                    break
        except asyncio.CancelledError:
            print("\n进度监听已中断")
        except Exception as e:
            print(f"\n进度监听出错: {str(e)}")

    async def chat_loop(self):
        """
        启动命令行交互式对话循环，支持文本查询和文件上传
        """
        print("\nMCP 客户端已启动，输入你的问题或'upload'上传文件，输入'quit'退出。")
        while True:
            try:
                query = input("\n请输入问题或命令: ").strip()
                if query.lower() == "quit":
                    break
                elif query.lower() == "upload":
                    await self.handle_file_upload()
                else:
                    response = await self.process_query(query)
                    print("\n" + response)
            except Exception as e:
                print(f"\n发生错误: {str(e)}")

async def main():
    """
    程序入口，设置配置并启动 MCP 客户端。
    """
    _ = load_dotenv(find_dotenv())
    my_api_key = os.environ['DPSEEK_API_KEY']
    # 从环境变量获取配置
    model_name = os.getenv("MODEL_NAME", "deepseek-chat")
    base_url = os.getenv("BASE_URL", "https://api.deepseek.com/v1")
    api_key = os.getenv("DPSEEK_API_KEY",my_api_key)

    # 支持工具调用的 llm
    # model_name = os.getenv("MODEL_NAME", "qwen2.5:32b")
    # base_url = os.getenv("BASE_URL", "http://192.168.0.184:11434/v1")
    # # api_key = os.getenv("OPENAI_API_KEY")
    # api_key = 'cannot be empty'

    # 不支持工具调用的 llm
    # model_name = os.getenv("MODEL_NAME", "huitian2")
    # base_url = os.getenv("BASE_URL", "http://192.168.0.184:9997/v1")
    # # api_key = os.getenv("OPENAI_API_KEY")
    # api_key = 'cannot be empty'

    if not api_key:
        print("未设置 API_KEY 环境变量。")
        sys.exit(1)

    # 定义 SSE 服务器地址列表
    server_urls = ["http://192.168.0.184:6001/sse"]

    # 创建并运行客户端
    client = MCPClient(model_name=model_name, base_url=base_url, api_key=api_key, server_urls=server_urls)
    try:
        await client.initialize_sessions()
        await client.chat_loop()
    finally:
        await client.cleanup()


if __name__ == "__main__":
    asyncio.run(main())
